package com.example.scribe;

// minimal imports
import java.util.HashMap;
import java.util.Map;

import com.example.scribe.error.ScribeError;

/**
 * Minimal type checker: for now only validates function call argument counts
 * against declared function parameter lists (happy-path).
 */
public class TypeChecker {
    private final Map<String, Integer> functionParamCounts = new HashMap<>();
    private final Map<String, java.util.List<String>> functionParamTypes = new HashMap<>();
    private final Map<String, String> functionReturnTypes = new HashMap<>();
    private final Map<String, String> variableTypes = new HashMap<>();

    public void check(SyntaxNode root) {
        // first pass: collect named function declarations
        collectFunctions(root);
        // second pass: validate calls
        validateCalls(root);
    }

    private void collectFunctions(SyntaxNode node) {
        if (node == null) return;
        if ("FunctionDecl".equals(node.type())) {
            String name = null;
            int paramCount = 0;
            java.util.List<String> ptypes = new java.util.ArrayList<>();
            String retType = null;
            for (SyntaxNode c : node.children()) {
                if ("Name".equals(c.type()) || "Identifier".equals(c.type())) {
                    name = c.text();
                }
                if ("Params".equals(c.type())) {
                    paramCount = c.children().size();
                    for (SyntaxNode p : c.children()) {
                        // each param node should have a Type child
                        String t = null;
                        for (SyntaxNode pc : p.children()) {
                            if ("Type".equals(pc.type())) t = pc.text();
                        }
                        ptypes.add(t);
                    }
                }
                if ("ReturnType".equals(c.type())) {
                    for (SyntaxNode rc : c.children()) if ("Type".equals(rc.type())) retType = rc.text();
                }
            }
            if (name != null) {
                functionParamCounts.put(name, paramCount);
                functionParamTypes.put(name, ptypes);
                if (retType != null) functionReturnTypes.put(name, retType);
            }
        }
        if ("VariableDecl".equals(node.type())) {
            String id = null;
            String t = null;
            for (SyntaxNode c : node.children()) {
                if ("Identifier".equals(c.type())) id = c.text();
                if ("Type".equals(c.type())) t = c.text();
            }
            if (id != null && t != null) variableTypes.put(id, t);
        }
        for (SyntaxNode c : node.children()) collectFunctions(c);
    }

    private void validateCalls(SyntaxNode node) {
        if (node == null) return;
        if ("Call".equals(node.type())) {
            // callee is expected as first child
            if (!node.children().isEmpty()) {
                SyntaxNode callee = node.children().get(0);
                String calleeName = extractIdentifierName(callee);
                if (calleeName != null && functionParamCounts.containsKey(calleeName)) {
                    int expected = functionParamCounts.get(calleeName);
                    int provided = 0;
                    java.util.List<String> providedTypes = new java.util.ArrayList<>();
                    for (SyntaxNode ch : node.children()) {
                        if ("Args".equals(ch.type())) {
                            provided = ch.children().size();
                            for (SyntaxNode a : ch.children()) providedTypes.add(inferType(a));
                            break;
                        }
                    }
                    if (provided != expected) {
                        throw ScribeError.typeError(node.text(), 1, 0,
                                "Function '" + calleeName + "' expects " + expected + " args but got " + provided);
                    }

                    // check types if we know expected types
                    java.util.List<String> expTypes = functionParamTypes.get(calleeName);
                    if (expTypes != null && !expTypes.isEmpty()) {
                        for (int i = 0; i < expTypes.size() && i < providedTypes.size(); i++) {
                            String exp = expTypes.get(i);
                            String prov = providedTypes.get(i);
                            if (exp != null && prov != null && !exp.equals(prov)) {
                                throw ScribeError.typeError(node.text(), 1, 0,
                                        "Argument type mismatch for function '" + calleeName + "' at pos " + i + ": expected " + exp + " but got " + prov);
                            }
                        }
                    }
                }
            }
        }
        for (SyntaxNode c : node.children()) validateCalls(c);
    }

    private String inferType(SyntaxNode expr) {
        if (expr == null) return null;
        // Terminal number
        if (expr.type() != null && expr.type().contains("NUMBER")) {
            String t = expr.text();
            if (t.matches(".*(i8|i16|i32|i64)$")) return t.replaceAll(".*(i8|i16|i32|i64)$", "$1");
            if (t.matches(".*(f32|f64)$")) return t.replaceAll(".*(f32|f64)$", "$1");
            return "i32"; // default integer
        }
        if (expr.type() != null && expr.type().contains("TRUE") || expr.type() != null && expr.type().contains("FALSE")) return "bool";
        // identifier
        if ("Identifier".equals(expr.type()) || "identifier".equals(expr.type()) || expr.type().equals("Terminal:IDENT")) {
            String name = expr.text();
            if (variableTypes.containsKey(name)) return variableTypes.get(name);
            return null;
        }
        // if node contains a child Type (e.g., cast nodes), prefer that
        for (SyntaxNode c : expr.children()) {
            if ("Type".equals(c.type())) return c.text();
        }
        // fallback: try to inspect child nodes
        for (SyntaxNode c : expr.children()) {
            String t = inferType(c);
            if (t != null) return t;
        }
        return null;
    }

    private String extractIdentifierName(SyntaxNode node) {
        if (node == null) return null;
        // common patterns: Name, identifier, Identifier
        if ("Name".equals(node.type()) || "Identifier".equals(node.type()) || "identifier".equals(node.type())) {
            return node.text();
        }
        // sometimes the identifier might be nested as child
        for (SyntaxNode c : node.children()) {
            String n = extractIdentifierName(c);
            if (n != null) return n;
        }
        return null;
    }
}
